[parquet] Add row group index virtual column by friendlymatthew · Pull Request #9117 · apache/arrow-rs

friendlymatthew · 2026-01-08T15:49:17Z

Which issue does this PR close?

Rationale for this change

This PR adds support for storing row group indices as a virtual column, allowing users to determine which row group each row originated from

The usage pattern is quite simple, something like:

use parquet::arrow::RowGroupIndex;

let row_group_index_field = Arc::new(
    Field::new("row_group_index", DataType::Int64, false)
        .with_extension_type(RowGroupIndex)
);

let options = ArrowReaderOptions::new()
    .with_virtual_columns(vec![row_group_index_field])?;

let reader = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options)?
    .build()?;

friendlymatthew · 2026-01-08T16:53:25Z

@Jefffrey if you have some bandwidth, I'd be curious to get your thoughts

friendlymatthew · 2026-01-09T14:58:32Z

cc @alamb

Jefffrey · 2026-01-09T15:00:29Z

@Jefffrey if you have some bandwidth, I'd be curious to get your thoughts

I'll see if I can take a look at this PR soon, though I will say it's been a while since I looked at the parquet codebase 😅

adriangb · 2026-01-09T16:10:49Z

This would be sweet! I like the idea of using an extension type as a marker, and letting the caller customize the column name.

alamb · 2026-01-09T17:26:37Z

I think the extension type marker was pioneered by @jkylling and @vustef

I will also try and review this sooner rather than later

scovich

Nice! (a few nits)

parquet/src/arrow/array_reader/row_group_index.rs

scovich · 2026-01-10T20:32:01Z

parquet/src/arrow/array_reader/row_group_index.rs

+    fn read_records(&mut self, batch_size: usize) -> Result<usize> {
+        let starting_len = self.buffered_indices.len();
+        self.buffered_indices
+            .extend((&mut self.remaining_indices).take(batch_size));


nit: I know this is how the row index reader did it, but since that code merged I learned that Iterator::by_ref is a thing.

Suggested change

.extend((&mut self.remaining_indices).take(batch_size));

.extend(self.remaining_indices.by_ref().take(batch_size));

It's not shorter, but does seem more readable?

(more below)

TIL. Thank you @scovich

scovich · 2026-01-10T20:38:47Z

parquet/src/arrow/schema/virtual_type.rs

+        if metadata.is_some_and(str::is_empty) {
+            Ok("")
+        } else {
+            Err(ArrowError::InvalidArgumentError(
+                "Virtual column extension type expects an empty string as metadata".to_owned(),
+            ))
+        }


nit: is a match simpler?

Suggested change

if metadata.is_some_and(str::is_empty) {

Ok("")

} else {

Err(ArrowError::InvalidArgumentError(

"Virtual column extension type expects an empty string as metadata".to_owned(),

))

}

match metadata {

Some(&"") => Ok(""),

_ => Err(ArrowError::InvalidArgumentError(

"Virtual column extension type expects an empty string as metadata".to_owned(),

)),

}

or even

Suggested change

if metadata.is_some_and(str::is_empty) {

Ok("")

} else {

Err(ArrowError::InvalidArgumentError(

"Virtual column extension type expects an empty string as metadata".to_owned(),

))

}

if let Some(&"") = metadata {

return Ok("");

};

Err(ArrowError::InvalidArgumentError(

"Virtual column extension type expects an empty string as metadata".to_owned(),

))

Jefffrey

Looks reasonable to me 👍

alamb

Thanks @friendlymatthew and @scovich and @Jefffrey -- this is really cool. I can't wait to get this integrated into downstream projects

alamb · 2026-01-14T22:40:09Z

parquet/src/arrow/array_reader/row_group_index.rs

+
+    fn consume_batch(&mut self) -> Result<ArrayRef> {
+        Ok(Arc::new(Int64Array::from_iter(
+            self.buffered_indices.drain(..),


Most of the time each ArrayReader will be instantiated for pages from exactly one row group we can probably make this significantly faster by optimizing the case when reading from only a single row group.

The other thing would be to pre-calculate the arrays (where the row group values don't change) and return the same ArrayRef until the RowGroup actually changes

However, there is no reason we need to do that now.

Filed:

perf: optimize RowGroupIndexReader for single row group reads #9180

perf: cache ArrayRef in RowGroupIndexReader when row group doesn't change #9181

I wonder if we should have the same optimization for RowNumberReader 🤔

alamb · 2026-01-14T22:42:26Z

parquet/src/arrow/array_reader/row_group_index.rs

+
+    fn skip_records(&mut self, num_records: usize) -> Result<usize> {
+        // TODO: Use advance_by when it stabilizes to improve performance
+        Ok((self.remaining_indices.by_ref()).take(num_records).count())


I don't see any test coverage for this method.

I ran code coverage

cargo llvm-cov --html --all-features -p parquet

And skip doesn't seem to be covered:

alamb

Thank you @friendlymatthew @Jefffrey and @scovich

github-actions bot added the parquet Changes to the parquet crate label Jan 8, 2026

friendlymatthew force-pushed the friendlymatthew/virtual-row-group-index branch 4 times, most recently from 073e3d5 to 0c9e12d Compare January 8, 2026 16:08

alamb mentioned this pull request Jan 9, 2026

Andrew Lamb Weekly-ish Open Source plan - 2026-01-05 apache/datafusion#19652

Closed

44 tasks

scovich approved these changes Jan 10, 2026

View reviewed changes

Jefffrey approved these changes Jan 13, 2026

View reviewed changes

friendlymatthew force-pushed the friendlymatthew/virtual-row-group-index branch from 0c9e12d to 0c019db Compare January 13, 2026 22:42

This was referenced Jan 14, 2026

[DISCUSSION] DataFusion Road Map: Q1 2026 apache/datafusion#18494

Open

Add input_file_name built-in function apache/datafusion#6051

Open

alamb reviewed Jan 14, 2026

View reviewed changes

This was referenced Jan 14, 2026

perf: optimize RowGroupIndexReader for single row group reads #9180

Closed

perf: cache ArrayRef in RowGroupIndexReader when row group doesn't change #9181

Open

Add row group index virtual column

4437dc6

friendlymatthew force-pushed the friendlymatthew/virtual-row-group-index branch from 0c019db to 4437dc6 Compare January 14, 2026 23:15

alamb approved these changes Jan 15, 2026

View reviewed changes

alamb merged commit d81d6c3 into apache:main Jan 15, 2026
16 checks passed

	.extend((&mut self.remaining_indices).take(batch_size));
	.extend(self.remaining_indices.by_ref().take(batch_size));

Conversation

friendlymatthew commented Jan 8, 2026

Which issue does this PR close?

Rationale for this change

Uh oh!

friendlymatthew commented Jan 8, 2026

Uh oh!

friendlymatthew commented Jan 9, 2026

Uh oh!

Jefffrey commented Jan 9, 2026

Uh oh!

adriangb commented Jan 9, 2026

Uh oh!

alamb commented Jan 9, 2026

Uh oh!

scovich left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

scovich Jan 10, 2026

Choose a reason for hiding this comment

Uh oh!

friendlymatthew Jan 13, 2026

Choose a reason for hiding this comment

Uh oh!

scovich Jan 10, 2026

Choose a reason for hiding this comment

Uh oh!

Jefffrey left a comment

Choose a reason for hiding this comment

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

alamb Jan 14, 2026

Choose a reason for hiding this comment

Uh oh!

friendlymatthew Jan 14, 2026

Choose a reason for hiding this comment

Uh oh!

alamb Jan 14, 2026

Choose a reason for hiding this comment

Uh oh!

alamb left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

alamb left a comment •

edited

Loading